// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/common/base/admission_controller.h>
#include <gtest/gtest.h>
#include <atomic>
#include <pollux/common/base/pollux_exception.h>
#include <pollux/testing/gtest_utils.h>

using namespace kumo::pollux;
namespace kumo::pollux::common {
TEST(AdmissionController, basic) {
  const uint64_t kLimit = 100000;
  AdmissionController::Config config;
  config.maxLimit = kLimit;
  AdmissionController admissionController(config);
  EXPECT_EQ(admissionController.currentResourceUsage(), 0);

  admissionController.accept(100);
  EXPECT_EQ(admissionController.currentResourceUsage(), 100);

  admissionController.accept(100);
  EXPECT_EQ(admissionController.currentResourceUsage(), 200);

  admissionController.release(100);
  EXPECT_EQ(admissionController.currentResourceUsage(), 100);

  POLLUX_ASSERT_THROW(
      admissionController.release(101),
      "Cannot release more units than have been acquired");

  POLLUX_ASSERT_THROW(
      admissionController.accept(kLimit + 1),
      "A single request cannot exceed the max limit");
}

TEST(AdmissionController, multiThreaded) {
  // Ensure that resource usage never exceeds the limit set in the admission
  // controller.
  const uint64_t kLimit = 30;
  std::atomic_uint64_t currentUsage{0};
  AdmissionController::Config config;
  config.maxLimit = kLimit;
  AdmissionController admissionController(config);

  std::vector<std::thread> threads;
  for (int i = 0; i < 20; i++) {
    threads.push_back(std::thread([&]() {
      for (int j = 0; j < 10'000; j++) {
        uint64_t usageUnits = std::rand() % (kLimit + 2);
        if (usageUnits > kLimit) {
          POLLUX_ASSERT_THROW(
              admissionController.accept(usageUnits),
              "A single request cannot exceed the max limit")
          continue;
        }
        admissionController.accept(usageUnits);
        uint64_t curr = currentUsage.fetch_add(usageUnits);
        ASSERT_LE(curr + usageUnits, kLimit);
        currentUsage.fetch_sub(usageUnits);
        admissionController.release(usageUnits);
      }
    }));
  }
  for (auto& thread : threads) {
    thread.join();
  }
}
} // namespace kumo::pollux::common
